import type { MarkdownRenderer } from 'vitepress'
import { StateInline, StateBlock } from 'markdown-it'

export default (md: MarkdownRenderer): void => {
  const annotation: Map<string, number> = new Map()
  md.renderer.rules.annotation_cit = (tokens, idx) => {
    const token = tokens[idx]
    const id = token.meta.id
    const label = token.meta.label
    let count: number
    if (annotation.has(id)) {
      count = annotation.get(id)!
    } else {
      const length = annotation.size
      count = length + 1
      annotation.set(id, count)
    }
    return `<span>${label}<sup class="annotation-ref"><a id="annotation-cit-${count}" href="#annotation-def-${count}" data-annotation-ref>${count}</a></sup></span>`
  }

  function annotation_cit(state: StateInline, silent: boolean) {
    const max = state.posMax
    const start = state.pos

    if (state.src.charCodeAt(start) !== 0x5b) return false // [
    if (state.src.charCodeAt(start + 1) !== 0x5e) return false // ^

    let pos = start + 2
    while (pos < max) {
      if (state.src.charCodeAt(pos) === 0x5d) break // ]
      pos++
    }

    if (pos === max) return false
    const label = state.src.slice(start + 2, pos)

    if (!silent) {
      const token = state.push('annotation_cit', '', 0)
      const safeId = `xy-${label.replace(/[^\w-]/g, '-')}`
      token.meta = { id: safeId, label }
    }

    state.pos = pos + 1
    return true
  }
  md.inline.ruler.before('emphasis', 'annotation_cit', annotation_cit)

  md.renderer.rules.annotation_def = (tokens, idx) => {
    const token = tokens[idx]
    const id = token.meta.id
    const content = token.meta.content
    const label = token.meta.label
    const count = annotation.get(id)!
    if (!count) return ''
    return `
          <div id="annotation-def-${count}" class="annotation-item">
            <span class="annotation-marker">${label}<sup class="annotation-ref"><a href="#annotation-cit-${count}" data-annotation-ref>${count}</a></sup>：</span>
            <span class="annotation-content">${md.renderInline(content)}</span>
            <a href="#annotation-cit-${count}" class="annotation-back-cit" aria-label="返回正文">↩</a>
          </div>
        `
  }

  function annotation_collect(state: StateBlock, startLine: number, _: number, silent: boolean) {
    if (silent) return true
    const line = state.src.slice(state.bMarks[startLine], state.eMarks[startLine])
    const match = line.match(/^\[\^([^\]]+)\]:\s*(.*)$/)
    if (!match) return false
    const token = new state.Token('annotation_def', '', 0)
    const label = match[1]
    const safeId = `xy-${label.replace(/[^\w-]/g, '-')}`
    token.meta = {
      id: safeId,
      label,
      content: match[2].trim()
    }
    token.block = true
    state.tokens.push(token)
    state.line = startLine + 1
    return true
  }

  md.block.ruler.before('reference', 'annotation_collect', annotation_collect, {
    alt: ['paragraph', 'reference']
  })
}
